# Testing radial basis functions.
# 23-10-26

import os
import sys
import jax
import jax.numpy as jnp
import numpy as np
import joblib
import haiku as hk
import mindspore as ms
import mindspore.nn as ms_nn

from jax import Array, jit
from typing import Optional, Union, List
from jax.nn.initializers import lecun_normal, normal, zeros, ones
from mindspore import load_param_into_net, Tensor, nn, context

sys.path.append(os.path.dirname(sys.path[0]))

# Test for utils
# setting context & hyper parameters
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")

from cybertron.utils.cutoff import CosineCutoff
from ms_cybertron.cutoff import CosineCutoff as ms_CosineCutoff

EPSILON = 1e-3
rng = jax.random.PRNGKey(42)

print("=================Test Cosine Cutoff====================")
test_input = np.random.uniform(low=0.2, high=1.2, size=(2, 3, 3))
test_mask = test_input < 0.8

def cos_fn(x, y, training=False):
    cos = CosineCutoff(cutoff=0.8)(x, y)
    return cos

cos_fn = jax.vmap(cos_fn, in_axes=(0, 0, ))
ts_cos_fn = hk.transform(cos_fn, apply_rng=True)
params = ts_cos_fn.init(rng, test_input, test_mask)
cos_out, cos_mask_out = jit(ts_cos_fn.apply)(params, rng, test_input, test_mask)
print(f"cos_out: shape of {cos_out.shape}")

ms_cos_fn = ms_CosineCutoff(cutoff=0.8)
ms_test_input = Tensor(test_input, ms.float32)
ms_test_mask = Tensor(test_mask, ms.bool_)
ms_cos_out, ms_mask_out = ms_cos_fn(ms_test_input, ms_test_mask)

print(f"ms_cos_out: shape of {ms_cos_out.shape}")
print("[from jax] cos_out: \n", cos_out)
print("[from mindspore] ms_cos_out: \n", ms_cos_out.asnumpy())
cos_same = np.allclose(cos_out, ms_cos_out.asnumpy(), atol=EPSILON)
print(f"cosine cutoff is the same: {cos_same} in accuracy of {EPSILON}.")
print(f"mask is the same: {np.all(cos_mask_out == ms_mask_out.asnumpy())}")

# caculate the norm of diff
cos_diff = cos_out - ms_cos_out.asnumpy()
cos_norm = np.linalg.norm(cos_diff)
print(f"norm of diff: {cos_norm}")

print("=================Test Smooth Cutoff====================")
from cybertron.utils.cutoff import SmoothCutoff
from ms_cybertron.cutoff import SmoothCutoff as ms_SmoothCutoff

test_input = np.random.uniform(low=0.2, high=1.2, size=(2, 3, 3))
test_mask = test_input < 0.8

def smooth_fn(x, y, training=False):
    smooth = SmoothCutoff(cutoff=0.8)(x, y)
    return smooth

smooth_fn = jax.vmap(smooth_fn, in_axes=(0, 0, ))
ts_smooth_fn = hk.transform(smooth_fn, apply_rng=True)
params = ts_smooth_fn.init(rng, test_input, test_mask)
smooth_out, smooth_mask_out = jit(ts_smooth_fn.apply)(params, rng, test_input, test_mask)
print(f"smooth_out: shape of {smooth_out.shape}")

ms_smooth_fn = ms_SmoothCutoff(cutoff=0.8)
ms_test_input = Tensor(test_input, ms.float32)
ms_test_mask = Tensor(test_mask, ms.bool_)
ms_smooth_out, ms_mask_out = ms_smooth_fn(ms_test_input, ms_test_mask)

print(f"ms_smooth_out: shape of {ms_smooth_out.shape}")
print("[from jax] smooth_out: \n", smooth_out)
print("[from mindspore] ms_smooth_out: \n", ms_smooth_out.asnumpy())
smooth_same = np.allclose(smooth_out, ms_smooth_out.asnumpy(), atol=EPSILON)
print(f"smooth cutoff is the same: {smooth_same} in accuracy of {EPSILON}.")
print(f"mask is the same: {np.all(smooth_mask_out == ms_mask_out.asnumpy())}")

# caculate the norm of diff
smooth_diff = smooth_out - ms_smooth_out.asnumpy()
smooth_norm = np.linalg.norm(smooth_diff)
print(f"norm of diff: {smooth_norm}")


